import torch
from torch import nn
from transformers import BertConfig, BertModel

from open_clip.model import CLIP

class CLIP_BERT(nn.Module):
    def __init__(self, clip, args):
        super().__init__()
        self.clip = clip

        # 文本后的线性层
        self.linear_cls = nn.Linear(1024, 512)

        """ 替换文本编码器 """
        config_bert = BertConfig.from_json_file(args.path_bert_model_config)
        self.bert = BertModel(config_bert)

        state_dict = torch.load(args.path_bert_state_dict)
        # 将权重赋值到模型中
        bert_static_dict = {}
        linear_static_dict = {}
        for key, value in state_dict.items():
            if 'transformer' in key:
                key = key.split(".", 1)[1]
                bert_static_dict[key] = value
            elif "Linear" in key:
                key = key.split(".", 1)[1]
                linear_static_dict[key] = value
            elif "embeddings" in key:
                bert_static_dict[key] = value

        bert_static_dict.pop("embeddings.position_ids")
        self.bert.load_state_dict(bert_static_dict)
        self.linear_cls.load_state_dict(linear_static_dict)


